package ffte.algorithm

/**
 * Prime number utilities for FFT algorithms
 *
 * This object provides mathematical utilities for prime number operations
 * that are essential for FFT algorithm decomposition and optimization.
 * Prime factorization is fundamental to FFT algorithms like PSB (Prime Space Base)
 * and determines the optimal decomposition strategies.
 *
 * Key Operations:
 * - Prime number generation and caching
 * - Prime factorization and representation
 * - GCD calculation for co-prime analysis
 * - Chinese Remainder Theorem for index mapping
 *
 * Usage in FFT Algorithms:
 * - PSB decomposition: Requires coprime factors for optimal performance
 * - Mixed-radix decomposition: Prime factorization determines decomposition strategy
 * - Winograd algorithms: Prime factors enable minimal multiplication algorithms
 */
object Prime {

    /**
     * Lazy stream of prime numbers
     *
     * Generates prime numbers on-demand using the Sieve of Eratosthenes
     * approach. The stream is lazy and infinite, generating primes as needed.
     *
     * Implementation: 2 :: Stream.from(3).filter(isPrimeByDivision)
     */
    lazy val primes: Stream[Int] = 2 #:: Stream.from(3).filter(i => primes.takeWhile(j => j * j <= i).forall(i % _ > 0))

    /**
     * Cached list of first 16384 prime numbers
     *
     * Provides fast access to commonly used prime numbers for
     * FFT decomposition. This cache avoids repeated prime
     * generation during algorithm processing.
     */
    val primeList = primes.take(16384).toList

    /**
     * Count prime factor p occurrences in integer i
     *
     * Recursively counts how many times prime p divides i and returns
     * the result along with the remaining quotient.
     *
     * @param i Integer to factor
     * @param n Current count (initially 0)
     * @param p Prime factor to count
     * @return Tuple of (quotient, count)
     */
    def split(i: Int, n: Int, p: Int): (Int, Int) = if (i % p == 0) split(i / p, n + 1, p) else (i, n)

    /**
     * Prime factorization with powers
     *
     * Factorizes integer N into its prime factors with their
     * corresponding powers. Returns a map where keys are primes
     * and values are the powers of each prime in the factorization.
     *
     * Example: N=60 → Map(2 → 2, 3 → 1, 5 → 1)
     *
     * @param N Integer to factorize
     * @return Map of prime factors to their powers
     */
    def repesent(N: Int) = (for (p <- primeList; if (N % p == 0)) yield {
        val (_, n) = split(N, 0, p)
        (p, n)
    }).toMap

    /**
     * Greatest Common Divisor calculation
     *
     * Implements Euclid's algorithm for finding the greatest common
     * divisor of two integers. This is essential for determining
     * whether two factors are coprime, which affects FFT algorithm
     * selection.
     *
     * @param a First integer
     * @param b Second integer
     * @return Greatest common divisor of a and b
     */
    def gcd(a: Int, b: Int): Int =
        if (b == 0) a else gcd(b, a % b)

    /**
     * Prime number checking
     *
     * Determines if a number is prime by checking against the
     * cached prime list. Fast for numbers up to the cached limit.
     *
     * @param a Number to check for primality
     * @return True if a is prime, false otherwise
     */
    def isPrime(a: Int): Boolean = primeList.contains(a)

    /**
     * Chinese Remainder Theorem base class
     *
     * Implements the mathematical foundation for index mapping in
     * PSB (Prime Space Base) FFT decomposition. The CRT enables optimal
     * reordering of data between FFT stages when factors are coprime.
     *
     * Mathematical Foundation:
     * For coprime integers m and n, the system:
     * x ≡ a (mod m)
     * x ≡ b (mod n)
     * has a unique solution modulo mn.
     *
     * Key Properties:
     * - Provides modular inverses for index calculations
     * - Enables efficient data reordering between FFT stages
     * - Essential for PSB algorithm optimization
     *
     * @param first First coprime factor
     * @param second Second coprime factor
     */
    case class Base(first: Int, second: Int) {
        /** Product of the two coprime factors */
        val N = first * second

        /** Modular inverse of second modulo first */
        val invsecond = BigInt(second).modInverse(BigInt(first)).toInt

        /** Modular inverse of first modulo second */
        val invfirst = second - (invsecond * invsecond - 1) / first

        /**
         * Calculate remainders modulo each factor
         *
         * @param n Number to compute remainders for
         * @return Tuple of (n mod first, n mod second)
         */
        def rem(n: Int) = (n % first, n % second)

        /**
         * Chinese Remainder Theorem computation
         *
         * Solves the CRT system:
         * x ≡ f (mod first)
         * x ≡ s (mod second)
         *
         * @param f Remainder modulo first
         * @param s Remainder modulo second
         * @return Solution x modulo (first*second)
         */
        def chinese(f: Int, s: Int) = (s * invfirst * first + f * invsecond * second) % N
    }
}